import os
import json

import numpy as np

if __name__ =='__main__':
    camera_calib_path = '/media/adt/ZWH4T/ZWH/Dataset3d/dataset_18xx-bag/cyw2.0/calib/lidar2img'
    lr_calib_path = '/media/adt/ZWH4T/ZWH/Dataset3d/dataset_18xx-bag/cyw2.0/calib/lr'
    json_path = '/media/adt/ZWH4T/ZWH/Dataset3d/dataset_18xx-bag/cyw2.0/calib/bags/calib.json'
    pos2camera = {
        'front':'camera75',
        'back':'camera77',
        'left':'camera81',
        'right':'camera80',
    }
    lr2pos = {
        'transform_lidar_f':'lidar3',
        'transform_lidar_b':'lidar4',
        'transform_lidar_l':'lidar2',
        'transform_lidar_r':'lidar1',
        'transform_radar_f':'radar3',
    }
    with open(lr_calib_path, 'r') as file:
        for line in file:
            if line.startswith('transform_'):
                # 读取接下来的4行，每行是一个矩阵行
                matrix_lines = []
                for _ in range(4):
                    matrix_line = next(file).strip().split()
                    # 将最后一个元素（通常是1，表示齐次坐标的w分量）也转换为float
                    matrix_line = [float(x) for x in matrix_line]
                    matrix_lines.append(matrix_line)
                lr2base_matrix = np.array(matrix_lines)

                with open(json_path, 'r') as file2:
                    calib_data = json.load(file2)

                calib_data[lr2pos[line.split('\n')[0]]]['transforms']['base_link'] = lr2base_matrix.tolist()

                with open(json_path, 'w') as f:
                    json.dump(calib_data, f, ensure_ascii=False, indent=2)

    for filename in sorted(os.listdir(camera_calib_path)):
        if filename not in pos2camera:
            continue
        camera_id = pos2camera[filename]
        filename = camera_calib_path + '/' + filename + '/base2'+camera_id

        # 用于存储读取的数据
        initial_ypr = None
        initial_xyz = None
        base2img_matrix = None

        # 打开文件并逐行读取
        with open(filename, 'r') as file:
            for line in file:
                # 跳过空行和注释行（假设注释行以 # 开头）
                if line.strip() == '' or line.startswith('#'):
                    continue

                    # 解析旋转角度（ypr）
                if line.startswith('旋转ypr'):
                    yaw = next(file).strip().split()[0]
                    pitch = next(file).strip().split()[0]
                    roll = next(file).strip().split()[0]
                    euler = np.array([float(yaw), float(pitch), float(roll)])
                    # euler = np.array([float(roll), float(pitch), float(yaw)])
                    from scipy.spatial.transform import Rotation as R
                    rotation_matrix = R.from_euler('ZYX', euler).as_matrix()

                    next(file).strip().split()
                    x = next(file).strip().split()[0]
                    y = next(file).strip().split()[0]
                    z = next(file).strip().split()[0]
                    position = np.array([float(x), float(y), float(z)])

                    base2camera = np.eye(4)
                    base2camera[:3, :3] = rotation_matrix
                    base2camera[:3, 3] = position
                    print(base2camera)
                    # img2base = np.linalg.inv(base2camera)
                    pass



                    # 解析变换矩阵（base2img）
                elif line.startswith('base2img:'):
                    # 读取接下来的4行，每行是一个矩阵行
                    matrix_lines = []
                    for _ in range(4):
                        matrix_line = next(file).strip().split()
                        # 将最后一个元素（通常是1，表示齐次坐标的w分量）也转换为float
                        matrix_line = [float(x) for x in matrix_line]
                        matrix_lines.append(matrix_line)
                    base2img_matrix = np.array(matrix_lines)
                    img2base = np.linalg.inv(base2img_matrix)

                    with open(json_path, 'r') as file2:
                        calib_data = json.load(file2)

                    calib_data[camera_id]['transforms']['base_link'] = img2base.tolist()

                    with open(json_path, 'w') as f:
                        json.dump(calib_data, f, ensure_ascii=False, indent=2)


                    # 如果已经读取了所有需要的数据，可以提前退出循环
                if initial_ypr is not None and initial_xyz is not None and base2img_matrix is not None:
                    break

                    # 输出读取的数据以验证
        print("Initial YPR:", initial_ypr)
        print("Initial XYZ:", initial_xyz)
        print("Base2Img Matrix:\n", base2img_matrix)
